// Copyright (C) 2024 Slack Technologies, LLC
// SPDX-License-Identifier: Apache-2.0
package com.slack.circuit.codegen

import com.google.devtools.ksp.processing.JvmPlatformInfo
import com.google.devtools.ksp.processing.PlatformInfo
import com.google.devtools.ksp.symbol.KSClassDeclaration
import com.google.devtools.ksp.symbol.KSDeclaration
import com.google.devtools.ksp.symbol.KSFunctionDeclaration
import com.slack.circuit.codegen.FactoryType.UI
import com.squareup.kotlinpoet.AnnotationSpec
import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.KModifier.ABSTRACT
import com.squareup.kotlinpoet.LambdaTypeName
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.TypeName
import com.squareup.kotlinpoet.TypeSpec

internal enum class CodegenMode {
  UNKNOWN {
    override fun supportsPlatforms(platforms: List<PlatformInfo>): Boolean {
      return false
    }

    override fun addInjectAnnotation(
      classBuilder: TypeSpec.Builder,
      constructorBuilder: FunSpec.Builder,
      options: CircuitOptions,
    ) {
      error("Unknown codegen mode, should not be called")
    }
  },
  /**
   * The Anvil Codegen mode
   *
   * This mode annotates generated factory types with `ContributesMultibinding`, allowing for Anvil
   * to automatically wire the generated class up to Dagger's multibinding system within a given
   * scope (e.g. AppScope).
   *
   * ```kotlin
   * @ContributesMultibinding(AppScope::class)
   * public class FavoritesPresenterFactory @Inject constructor(
   *   private val factory: FavoritesPresenter.Factory,
   * ) : Presenter.Factory { ... }
   * ```
   */
  ANVIL {
    private val contributesMultibindingCN =
      ClassName("com.squareup.anvil.annotations", "ContributesMultibinding")

    override fun supportsPlatforms(platforms: List<PlatformInfo>): Boolean {
      // Anvil only supports JVM & Android
      return platforms.all { it is JvmPlatformInfo }
    }

    override fun annotateFactory(builder: TypeSpec.Builder, scope: TypeName) {
      builder.addAnnotation(
        AnnotationSpec.builder(contributesMultibindingCN).addMember("%T::class", scope).build()
      )
    }

    override fun addInjectAnnotation(
      classBuilder: TypeSpec.Builder,
      constructorBuilder: FunSpec.Builder,
      options: CircuitOptions,
    ) {
      constructorBuilder.addAnnotation(runtime.inject(options))
    }
  },

  /**
   * The Hilt Codegen mode
   *
   * This mode provides an additional type, a Hilt module, which binds the generated factory, wiring
   * up multibinding in the Hilt DI framework. The scope provided via `CircuitInject` is used to
   * define the dagger component the factory provider is installed in.
   *
   * ```kotlin
   * @Module
   * @InstallIn(SingletonComponent::class)
   * @OriginatingElement(topLevelClass = FavoritesPresenter::class)
   * public abstract class FavoritesPresenterFactoryModule {
   *   @Binds
   *   @IntoSet
   *   public abstract
   *       fun bindFavoritesPresenterFactory(favoritesPresenterFactory: FavoritesPresenterFactory):
   *       Presenter.Factory
   * }
   * ```
   */
  HILT {
    override val originAnnotation: OriginAnnotation = CircuitNames.DAGGER_ORIGIN

    override fun supportsPlatforms(platforms: List<PlatformInfo>): Boolean {
      // Hilt only supports JVM & Android
      return platforms.all { it is JvmPlatformInfo }
    }

    override fun produceAdditionalTypeSpec(
      factory: ClassName,
      factoryType: FactoryType,
      scope: TypeName,
      topLevelClass: ClassName?,
    ): TypeSpec {
      val moduleAnnotations =
        listOfNotNull(
          AnnotationSpec.builder(CircuitNames.DAGGER_MODULE).build(),
          AnnotationSpec.builder(CircuitNames.DAGGER_INSTALL_IN)
            .addMember("%T::class", scope)
            .build(),
          // Required to generate this here too since this is a separate class from the generated
          // factory
          topLevelClass?.let {
            AnnotationSpec.builder(CircuitNames.DAGGER_ORIGINATING_ELEMENT)
              .addMember("%L = %T::class", "topLevelClass", topLevelClass)
              .build()
          },
        )

      val providerAnnotations =
        listOf(
          AnnotationSpec.builder(CircuitNames.DAGGER_BINDS).build(),
          AnnotationSpec.builder(CircuitNames.DAGGER_INTO_SET).build(),
        )

      val providerReturns =
        if (factoryType == UI) {
          CircuitNames.CIRCUIT_UI_FACTORY
        } else {
          CircuitNames.CIRCUIT_PRESENTER_FACTORY
        }

      val factoryName = factory.simpleName

      val providerSpec =
        FunSpec.builder("bind${factoryName}")
          .addModifiers(ABSTRACT)
          .addAnnotations(providerAnnotations)
          .addParameter(name = factoryName.replaceFirstChar { it.lowercase() }, type = factory)
          .returns(providerReturns)
          .build()

      return TypeSpec.classBuilder(factory.peerClass(factoryName + CircuitNames.MODULE))
        .addModifiers(ABSTRACT)
        .addAnnotations(moduleAnnotations)
        .addFunction(providerSpec)
        .build()
    }

    override fun addInjectAnnotation(
      classBuilder: TypeSpec.Builder,
      constructorBuilder: FunSpec.Builder,
      options: CircuitOptions,
    ) {
      constructorBuilder.addAnnotation(runtime.inject(options))
    }
  },

  /**
   * The `kotlin-inject` Anvil Codegen mode
   *
   * This mode annotates generated factory types with `ContributesMultibinding`, allowing for
   * KI-Anvil to automatically wire the generated class up to KI's multibinding system within a
   * given scope (e.g. AppScope).
   *
   * ```kotlin
   * @Inject
   * @ContributesMultibinding(AppScope::class)
   * public class FavoritesPresenterFactory(
   *   private val provider: () -> FavoritesPresenter,
   * ) : Presenter.Factory { ... }
   * ```
   */
  KOTLIN_INJECT_ANVIL {
    override val runtime: InjectionRuntime = InjectionRuntime.KotlinInject
    override val originAnnotation: OriginAnnotation = CircuitNames.KotlinInject.Anvil.ORIGIN

    override fun supportsPlatforms(platforms: List<PlatformInfo>): Boolean {
      // KI-Anvil supports all
      return true
    }

    override fun annotateFactory(builder: TypeSpec.Builder, scope: TypeName) {
      builder.addAnnotation(
        AnnotationSpec.builder(CircuitNames.KotlinInject.Anvil.CONTRIBUTES_BINDING)
          .addMember("%T::class", scope)
          .addMember("multibinding = true")
          .build()
      )
    }

    override fun addInjectAnnotation(
      classBuilder: TypeSpec.Builder,
      constructorBuilder: FunSpec.Builder,
      options: CircuitOptions,
    ) {
      classBuilder.addAnnotation(runtime.inject(options))
    }

    override fun filterValidInjectionSites(
      candidates: Collection<KSDeclaration>
    ): Collection<KSDeclaration> {
      return candidates.filter { it is KSFunctionDeclaration || it is KSClassDeclaration }
    }
  },

  /**
   * The `metro` code gen mode
   *
   * This mode annotates generated factory types with `ContributesIntoSet`, allowing for KI-Anvil to
   * automatically wire the generated class up to KI's multibinding system within a given scope
   * (e.g. AppScope).
   *
   * ```kotlin
   * @Inject
   * @ContributesIntoSet(AppScope::class)
   * public class FavoritesPresenterFactory(
   *   private val provider: Provider<FavoritesPresenter>
   * ) : Presenter.Factory { ... }
   * ```
   */
  METRO {
    override val runtime: InjectionRuntime = InjectionRuntime.Metro
    override val originAnnotation: OriginAnnotation = CircuitNames.Metro.ORIGIN

    override fun supportsPlatforms(platforms: List<PlatformInfo>): Boolean {
      // Metro supports all
      return true
    }

    override fun annotateFactory(builder: TypeSpec.Builder, scope: TypeName) {
      builder.addAnnotation(
        AnnotationSpec.builder(CircuitNames.Metro.CONTRIBUTES_INTO_SET)
          .addMember("%T::class", scope)
          .build()
      )
    }

    override fun addInjectAnnotation(
      classBuilder: TypeSpec.Builder,
      constructorBuilder: FunSpec.Builder,
      options: CircuitOptions,
    ) {
      classBuilder.addAnnotation(runtime.inject(options))
    }

    override fun filterValidInjectionSites(
      candidates: Collection<KSDeclaration>
    ): Collection<KSDeclaration> {
      return candidates.filter { it is KSFunctionDeclaration || it is KSClassDeclaration }
    }
  };

  open val runtime: InjectionRuntime = InjectionRuntime.Jakarta
  open val originAnnotation: OriginAnnotation? = null

  open fun annotateFactory(builder: TypeSpec.Builder, scope: TypeName) {}

  open fun produceAdditionalTypeSpec(
    factory: ClassName,
    factoryType: FactoryType,
    scope: TypeName,
    topLevelClass: ClassName?,
  ): TypeSpec? {
    return null
  }

  abstract fun supportsPlatforms(platforms: List<PlatformInfo>): Boolean

  abstract fun addInjectAnnotation(
    classBuilder: TypeSpec.Builder,
    constructorBuilder: FunSpec.Builder,
    options: CircuitOptions,
  )

  /** Filters the candidates for @Inject annotation placement. */
  open fun filterValidInjectionSites(
    candidates: Collection<KSDeclaration>
  ): Collection<KSDeclaration> = candidates.filterIsInstance<KSFunctionDeclaration>()

  sealed interface InjectionRuntime {

    fun inject(options: CircuitOptions): ClassName

    fun declarationInjects(options: CircuitOptions): Collection<ClassName> = listOf(inject(options))

    val assisted: ClassName
    val assistedInject: ClassName?
    val assistedFactory: ClassName?

    fun asProvider(providedType: TypeName, options: CircuitOptions): TypeName

    fun getProviderBlock(provider: CodeBlock): CodeBlock

    data object Jakarta : InjectionRuntime {
      override fun inject(options: CircuitOptions): ClassName =
        if (options.useJavaxOnly) CircuitNames.INJECT_JAVAX else CircuitNames.INJECT

      override fun declarationInjects(options: CircuitOptions): Collection<ClassName> {
        // If explicitly using javax only look for javax, otherwise look for both.
        return if (options.useJavaxOnly) {
          listOf(CircuitNames.INJECT_JAVAX)
        } else {
          listOf(CircuitNames.INJECT, CircuitNames.INJECT_JAVAX)
        }
      }

      override val assisted: ClassName = CircuitNames.ASSISTED
      override val assistedInject: ClassName = CircuitNames.ASSISTED_INJECT
      override val assistedFactory: ClassName = CircuitNames.ASSISTED_FACTORY

      override fun asProvider(providedType: TypeName, options: CircuitOptions): TypeName {
        val className =
          if (options.useJavaxOnly) CircuitNames.PROVIDER_JAVAX else CircuitNames.PROVIDER
        return className.parameterizedBy(providedType)
      }

      override fun getProviderBlock(provider: CodeBlock): CodeBlock {
        return CodeBlock.of("%L.get()", provider)
      }
    }

    data object KotlinInject : InjectionRuntime {
      override fun inject(options: CircuitOptions) = CircuitNames.KotlinInject.INJECT

      override val assisted: ClassName = CircuitNames.KotlinInject.ASSISTED
      override val assistedInject: ClassName? = null
      override val assistedFactory: ClassName? = null // It has no annotation

      override fun asProvider(providedType: TypeName, options: CircuitOptions): TypeName {
        return LambdaTypeName.get(returnType = providedType)
      }

      override fun getProviderBlock(provider: CodeBlock): CodeBlock {
        return CodeBlock.of("%L()", provider)
      }
    }

    data object Metro : InjectionRuntime {
      override fun inject(options: CircuitOptions) = CircuitNames.Metro.INJECT

      override val assisted: ClassName = CircuitNames.Metro.ASSISTED
      override val assistedInject: ClassName? = null
      override val assistedFactory: ClassName = CircuitNames.Metro.ASSISTED_FACTORY

      override fun asProvider(providedType: TypeName, options: CircuitOptions): TypeName {
        return CircuitNames.Metro.PROVIDER.parameterizedBy(providedType)
      }

      override fun getProviderBlock(provider: CodeBlock): CodeBlock {
        return CodeBlock.of("%L()", provider)
      }
    }
  }
}

public class OriginAnnotation(
  public val className: ClassName,
  public val parameterName: String? = null,
)
